CloudFormation一撃で作るKnowledge Base for Amazon Bedrock (with Aurora)

こんにちは、つくぼし(tsukuboshi0755)です！

先日以下のアップデートにより、Knowledge Base for Amazon BedrockとAgents for Amazon BedrockがCloudFormationによるデプロイをサポートしました。

今回はこのアップデートを利用して、ベクターストアにAuroraを利用したKnowledge Base for Amazon Bedrockを、CloudFormation一撃で作成可能なテンプレートについて紹介します！

なおベクターストアをセットアップするために一部カスタムリソースを含む形になるため、その点にご留意頂けるとありがたいです。

構成

今回CloudFormationで作成する構成図は以下のとおりです。

Knowledge Base for Amazon Bedrockに必要なリソースとして、以下のリソースも合わせて作成します。

• ベクターストア：Aurora PostgreSQL
• 機密情報：Secrets Manager
• データソース：S3

またLambda-Backed Custom Resourceを利用して、Auroraに対するデータベースのセットアップを自動化します。

テンプレート

全体のコードは以下の通りです。

Description: 'Knowledge Base for Amazon Bedrock with Aurora PostgreSQL (including db setup)'

Mappings:
  DatabaseMap:
    DatabaseName:
      Name: bedrockkbdb
    TableName:
      Name: bedrock_integration.bedrock_kb
    SchemaName:
      Name: bedrock_integration
    Username:
      Name: bedrock_user
    PrimaryKeyField:
      Name: id
    VectorField:
      Name: embedding
    TextField:
      Name: chunks
    MetadataField:
      Name: metadata

Parameters:
  DatabasePassword:
    Type: String
    Default: 'P@ssword123'
    Description: 'The password for the database user.'
    NoEcho: true
  EmbeddingModelId:
    Type: String
    Default: amazon.titan-embed-text-v1
    AllowedValues:
      - amazon.titan-embed-text-v1
      - cohere.embed-multilingual-v3
      - cohere.embed-english-v3
    Description: 'The Id of the Bedrock model that is used to generate embeddings.'

Conditions:
  IsTitanEmbedTextV1: !Equals [!Ref EmbeddingModelId, "amazon.titan-embed-text-v1"]

Resources:
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 192.168.0.0/24
      EnableDnsSupport: true
      EnableDnsHostnames: true
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-vpc
  PrivateSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Sub ${AWS::Region}a
      CidrBlock: 192.168.0.0/28
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-1a
  PrivateSubnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Sub ${AWS::Region}c
      CidrBlock: 192.168.0.16/28
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-1c
  PrivateRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-rtb
  PrivateSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      SubnetId: !Ref PrivateSubnet1
  PrivateSubnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      SubnetId: !Ref PrivateSubnet2
  DBSubnetGroup:
    Type: AWS::RDS::DBSubnetGroup
    Properties:
      DBSubnetGroupDescription: Bedrock RDS Subnet Group
      SubnetIds:
        - !Ref PrivateSubnet1
        - !Ref PrivateSubnet2
  AuroraCluster:
    Type: AWS::RDS::DBCluster
    DeletionPolicy: Delete
    Properties:
      DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
      Engine: aurora-postgresql
      EngineVersion: 15.5
      DBSubnetGroupName: !Ref DBSubnetGroup
      MasterUsername: postgresql
      ManageMasterUserPassword: true
      ServerlessV2ScalingConfiguration:
        MinCapacity: 0.5
        MaxCapacity: 1.0
      StorageEncrypted: true
      EnableHttpEndpoint: true
  AuroraDBInstance:
    Type: AWS::RDS::DBInstance
    DeletionPolicy: Delete
    Properties:
      Engine: aurora-postgresql
      DBInstanceClass: db.serverless
      DBClusterIdentifier: !Ref AuroraCluster
  SecretForAurora:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub ${AWS::StackName}-db-secret-for-bedrock
      Description: 'Secret for the database user for Bedrock'
      SecretString: !Sub
        - '{ "username":"${DatabaseUser}", "password":"${DatabasePassword}"}'
        - DatabaseUser: !FindInMap [DatabaseMap, Username, Name]
  DataSourceBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
            BucketKeyEnabled: true
      OwnershipControls:
        Rules:
          - ObjectOwnership: BucketOwnerEnforced
  BedrockAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 'bedrock:InvokeModel'
            Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      Roles:
        - !Ref BedrockKnowledgeBaseRole
  SecretsAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-secret-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 'secretsmanager:GetSecretValue'
            Resource: !Ref SecretForAurora
      Roles:
        - !Ref BedrockKnowledgeBaseRole
  AuroraAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-aurora-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 'rds:DescribeDBClusters'
              - 'rds-data:BatchExecuteStatement'
              - 'rds-data:ExecuteStatement'
            Resource: !GetAtt AuroraCluster.DBClusterArn
          - Effect: Allow
            Action:
              - 'secretsmanager:GetSecretValue'
            Resource: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
      Roles:
        - !Ref BedrockKnowledgeBaseRole
        - !Ref ExecSQLFunctionRole
  S3AccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-s3-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: S3ListBucketStatement
            Effect: Allow
            Action:
              - 's3:ListBucket'
            Resource: !GetAtt DataSourceBucket.Arn
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
          - Sid: S3GetObjectStatement
            Effect: Allow
            Action:
              - 's3:GetObject'
            Resource: !Sub
              - '${DataSourceBucketArn}/*'
              - DataSourceBucketArn: !GetAtt DataSourceBucket.Arn
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
      Roles:
        - !Ref BedrockKnowledgeBaseRole
  SetupAuroraData:
    Type: 'Custom::SetupAuroraData'
    Properties:
      ServiceToken: !GetAtt ExecSQLFunction.Arn
      Dimension: !If [IsTitanEmbedTextV1, 1536, 1024]
      ResourceArn: !GetAtt AuroraCluster.DBClusterArn
      SecretArn: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
      DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
      DatabasePassword: !Ref DatabasePassword
      TableName: !FindInMap [DatabaseMap, TableName, Name]
      SchemaName: !FindInMap [DatabaseMap, SchemaName, Name]
      UserName: !FindInMap [DatabaseMap, Username, Name]
      MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
      PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
      TextField: !FindInMap [DatabaseMap, TextField, Name]
      VectorField: !FindInMap [DatabaseMap, VectorField, Name]
    DependsOn: AuroraDBInstance
  ExecSQLFunctionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Sub ${AWS::StackName}-execsql-function-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  ExecSQLFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub ${AWS::StackName}-execsql-function
      Handler: index.lambda_handler
      Role: !GetAtt ExecSQLFunctionRole.Arn
      Runtime: python3.12
      Timeout: 600
      LoggingConfig:
        LogFormat: JSON
        ApplicationLogLevel: INFO
        SystemLogLevel: INFO
      Code:
        ZipFile: |
          import boto3
          import cfnresponse
          import logging
          from typing import Any, Dict

          logger = logging.getLogger()

          rds_data = boto3.client('rds-data')

          def execute_statement(resource_arn: str, database_name: str, secret_arn: str, sql: str) -> Any:
              response = rds_data.execute_statement(
                  resourceArn=resource_arn,
                  database=database_name,
                  secretArn=secret_arn,
                  sql=sql
              )
              return response


          def lambda_handler(event: Dict[str, Any], context: Any) -> None:
              try:
                  dimension = event['ResourceProperties']['Dimension']
                  resource_arn = event['ResourceProperties']['ResourceArn']
                  secret_arn = event['ResourceProperties']['SecretArn']
                  database_name = event['ResourceProperties']['DatabaseName']
                  database_password = event['ResourceProperties']['DatabasePassword']
                  table_name = event['ResourceProperties']['TableName']
                  schema_name = event['ResourceProperties']['SchemaName']
                  user_name = event['ResourceProperties']['UserName']
                  metadata_field = event['ResourceProperties']['MetadataField']
                  primary_key_field = event['ResourceProperties']['PrimaryKeyField']
                  text_field = event['ResourceProperties']['TextField']
                  vector_field = event['ResourceProperties']['VectorField']

                  if event['RequestType'] == 'Create':
                      create_extension = f"""
                      CREATE EXTENSION IF NOT EXISTS vector;
                      """
                      create_extension_res = execute_statement(resource_arn, database_name, secret_arn, create_extension)
                      logger.info(f"Create Extension Response: {create_extension_res}")

                      create_role = f"""
                      CREATE ROLE {user_name} WITH PASSWORD '{database_password}' LOGIN;
                      """
                      create_role_res = execute_statement(resource_arn, database_name, secret_arn, create_role)
                      logger.info(f"Create Role Response: {create_role_res}")

                      create_schema = f"""
                      CREATE SCHEMA {schema_name};
                      """
                      create_schema_res = execute_statement(resource_arn, database_name, secret_arn, create_schema)
                      logger.info(f"Create Schema Response: {create_schema_res}")

                      grant_schema = f"""
                      GRANT ALL ON SCHEMA {schema_name} to {user_name};
                      """
                      grant_schema_res = execute_statement(resource_arn, database_name, secret_arn, grant_schema)
                      logger.info(f"Grant Schema Response: {grant_schema_res}")

                      create_table = f"""
                      CREATE TABLE {table_name} ({primary_key_field} uuid PRIMARY KEY, {vector_field} vector({dimension}), {text_field} text, {metadata_field} json)
                      """
                      create_table_res = execute_statement(resource_arn, database_name, secret_arn, create_table)
                      logger.info(f"Create Table Response: {create_table_res}")

                      grant_table = f"""
                      GRANT ALL ON TABLE {table_name} TO {user_name};
                      """
                      grant_table_res = execute_statement(resource_arn, database_name, secret_arn, grant_table)
                      logger.info(f"Grant Table Response: {grant_table_res}")

                      create_index = f"""
                      CREATE INDEX on {table_name} USING hnsw ({vector_field} vector_cosine_ops);
                      """
                      create_index_res = execute_statement(resource_arn, database_name, secret_arn, create_index)
                      logger.info(f"Create Index Response: {create_index_res}")

                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
                  if event['RequestType'] == 'Update':
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
                  if event['RequestType'] == 'Delete':
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as e:
                  cfnresponse.send(event, context, cfnresponse.FAILED, {'Message': str(e)})
  BedrockKnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [bedrock.amazonaws.com]
            Action: ['sts:AssumeRole']
  BedrockKnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      Name: !Sub ${AWS::StackName}-knowledge-base
      KnowledgeBaseConfiguration:
          Type: VECTOR
          VectorKnowledgeBaseConfiguration:
              EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
      StorageConfiguration:
        Type: RDS
        RdsConfiguration:
          CredentialsSecretArn: !Ref SecretForAurora
          DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
          FieldMapping:
            MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
            PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
            TextField: !FindInMap [DatabaseMap, TextField, Name]
            VectorField: !FindInMap [DatabaseMap, VectorField, Name]
          ResourceArn: !GetAtt AuroraCluster.DBClusterArn
          TableName: !FindInMap [DatabaseMap, TableName, Name]
    DependsOn: SetupAuroraData
  BedrockKnowledgeBaseDS:
    Type: AWS::Bedrock::DataSource
    Properties:
      KnowledgeBaseId: !Ref BedrockKnowledgeBase
      Name: !Sub ${AWS::StackName}-data-source
      DataSourceConfiguration:
        Type: S3
        S3Configuration:
          BucketArn: !GetAtt DataSourceBucket.Arn

Outputs:
  BedrockKnowledgeBaseId:
    Value: !Ref BedrockKnowledgeBase
  BedrockDataSourceId:
    Value: !Ref BedrockKnowledgeBaseDS
  AuroraClusterId:
    Value: !Ref AuroraCluster
  DSBucketName:
    Value: !Ref DataSourceBucket

またSAMコードについても、以下のリポジトリに格納していますので、必要に応じてご参照ください。

以下では、テンプレートの主要な部分を抜粋して説明します。

VPC

  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 192.168.0.0/24
      EnableDnsSupport: true
      EnableDnsHostnames: true
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-vpc
  PrivateSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Sub ${AWS::Region}a
      CidrBlock: 192.168.0.0/28
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-1a
  PrivateSubnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Sub ${AWS::Region}c
      CidrBlock: 192.168.0.16/28
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-1c
  PrivateRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-private-subnet-rtb
  PrivateSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      SubnetId: !Ref PrivateSubnet1
  PrivateSubnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      SubnetId: !Ref PrivateSubnet2
  DBSubnetGroup:
    Type: AWS::RDS::DBSubnetGroup
    Properties:
      DBSubnetGroupDescription: Bedrock RDS Subnet Group
      SubnetIds:
        - !Ref PrivateSubnet1
        - !Ref PrivateSubnet2

Auroraに必要なVPC、サブネット、ルートテーブルを作成します。

なお今回の主題ではないため、詳細な説明は割愛します。

Aurora

  DBSubnetGroup:
    Type: AWS::RDS::DBSubnetGroup
    Properties:
      DBSubnetGroupDescription: Bedrock RDS Subnet Group
      SubnetIds:
        - !Ref PrivateSubnet1
        - !Ref PrivateSubnet2
  AuroraCluster:
    Type: AWS::RDS::DBCluster
    DeletionPolicy: Delete
    Properties:
      DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
      Engine: aurora-postgresql
      EngineVersion: 15.5
      DBSubnetGroupName: !Ref DBSubnetGroup
      MasterUsername: postgresql
      ManageMasterUserPassword: true
      ServerlessV2ScalingConfiguration:
        MinCapacity: 0.5
        MaxCapacity: 1.0
      StorageEncrypted: true
      EnableHttpEndpoint: true
  AuroraDBInstance:
    Type: AWS::RDS::DBInstance
    DeletionPolicy: Delete
    Properties:
      Engine: aurora-postgresql
      DBInstanceClass: db.serverless
      DBClusterIdentifier: !Ref AuroraCluster

Knowledge Base for Amazon Bedrockのベクターストアとして必要なAurora PostgreSQLを作成します。

MasterUsername及びManageMasterUserPasswordの両方を指定する事で、以下の通りAdminユーザー用Secrets ManagerがAuroraと同時に作成されます。

またMinCapacity及びMaxCapacityには、ACUのスケーリング設定を指定します。

さらにEnableHttpEndpointを指定する事で、後述のLambda-Backed Custom Resourceで使用するRDS Data APIを有効化します。

なお今回は以下を参考に、検証用途のためバックアップ不要としてDeletionPolicyをDeleteで指定していますが、必要に応じてSnapshotやRetainに変更してください。

Secrets Manager

  SecretForAurora:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub ${AWS::StackName}-db-secret-for-bedrock
      Description: 'Secret for the database user for Bedrock'
      SecretString: !Sub
        - '{ "username":"${DatabaseUser}", "password":"${DatabasePassword}"}'
        - DatabaseUser: !FindInMap [DatabaseMap, Username, Name]

Knowledge Base for Amazon BedrockからAuroraに接続するためのユーザー/パスワード情報をSecrets Managerから取得する必要があるため、Secrets Managerを作成します。

なお本来認証情報をハードコーディングする事はセキュリティ上好ましくない事にご留意ください。

S3 Bucket

  DataSourceBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
            BucketKeyEnabled: true
      OwnershipControls:
        Rules:
          - ObjectOwnership: BucketOwnerEnforced

Knowledge Base for Amazon Bedrockのデータソースとして必要なS3バケットを作成します。

このバケットに対してユーザーがドキュメントをアップロードする事で、Knowledge Base for Amazon Bedrockがデータソースとして利用できるようになります。

IAM Policy

  BedrockAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 'bedrock:InvokeModel'
            Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      Roles:
        - !Ref BedrockKnowledgeBaseRole
  SecretsAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-secret-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - 'secretsmanager:GetSecretValue'
            Resource: !Ref SecretForAurora
      Roles:
        - !Ref BedrockKnowledgeBaseRole
  AuroraAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-aurora-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - rds:DescribeDBClusters
              - rds-data:BatchExecuteStatement
              - rds-data:ExecuteStatement
            Resource: !GetAtt AuroraCluster.DBClusterArn
          - Effect: Allow
            Action:
              - 'secretsmanager:GetSecretValue'
            Resource: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
      Roles:
        - !Ref BedrockKnowledgeBaseRole
        - !Ref ExecSQLFunctionRole
  S3AccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-s3-access-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: S3ListBucketStatement
            Effect: Allow
            Action:
              - 's3:ListBucket'
            Resource: !GetAtt DataSourceBucket.Arn
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
          - Sid: S3GetObjectStatement
            Effect: Allow
            Action:
              - 's3:GetObject'
            Resource: !Sub
              - '${DataSourceBucketArn}/*'
              - DataSourceBucketArn: !GetAtt DataSourceBucket.Arn
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
      Roles:
        - !Ref BedrockKnowledgeBaseRole

Lambda-Backed Custom Resource及びKnowledge Base for Amazon Bedrockに必要なIAMポリシーを作成します。

Bedrock、Secrets Manager、Aurora、S3に対するアクセス権限を設定し、IAMロールにアタッチします。

Lambda-Backed Custom Resource

  SetupAuroraData:
    Type: 'Custom::SetupAuroraData'
    Properties:
      ServiceToken: !GetAtt ExecSQLFunction.Arn
      Dimension: !If [IsTitanEmbedTextV1, 1536, 1024]
      ResourceArn: !GetAtt AuroraCluster.DBClusterArn
      SecretArn: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
      DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
      DatabasePassword: !Ref DatabasePassword
      TableName: !FindInMap [DatabaseMap, TableName, Name]
      SchemaName: !FindInMap [DatabaseMap, SchemaName, Name]
      UserName: !FindInMap [DatabaseMap, Username, Name]
      MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
      PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
      TextField: !FindInMap [DatabaseMap, TextField, Name]
      VectorField: !FindInMap [DatabaseMap, VectorField, Name]
    DependsOn: AuroraDBInstance
  ExecSQLFunctionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Sub ${AWS::StackName}-execsql-function-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  ExecSQLFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub ${AWS::StackName}-execsql-function
      Handler: index.lambda_handler
      Role: !GetAtt ExecSQLFunctionRole.Arn
      Runtime: python3.12
      Timeout: 600
      LoggingConfig:
        LogFormat: JSON
        ApplicationLogLevel: INFO
        SystemLogLevel: INFO
      Code:
        ZipFile: |
          import boto3
          import cfnresponse
          import logging
          from typing import Any, Dict

          logger = logging.getLogger()

          rds_data = boto3.client('rds-data')

          def execute_statement(resource_arn: str, database_name: str, secret_arn: str, sql: str) -> Any:
              response = rds_data.execute_statement(
                  resourceArn=resource_arn,
                  database=database_name,
                  secretArn=secret_arn,
                  sql=sql
              )
              return response


          def lambda_handler(event: Dict[str, Any], context: Any) -> None:
              try:
                  dimension = event['ResourceProperties']['Dimension']
                  resource_arn = event['ResourceProperties']['ResourceArn']
                  secret_arn = event['ResourceProperties']['SecretArn']
                  database_name = event['ResourceProperties']['DatabaseName']
                  database_password = event['ResourceProperties']['DatabasePassword']
                  table_name = event['ResourceProperties']['TableName']
                  schema_name = event['ResourceProperties']['SchemaName']
                  user_name = event['ResourceProperties']['UserName']
                  metadata_field = event['ResourceProperties']['MetadataField']
                  primary_key_field = event['ResourceProperties']['PrimaryKeyField']
                  text_field = event['ResourceProperties']['TextField']
                  vector_field = event['ResourceProperties']['VectorField']

                  if event['RequestType'] == 'Create':
                      create_extension = f"""
                      CREATE EXTENSION IF NOT EXISTS vector;
                      """
                      create_extension_res = execute_statement(resource_arn, database_name, secret_arn, create_extension)
                      logger.info(f"Create Extension Response: {create_extension_res}")

                      create_role = f"""
                      CREATE ROLE {user_name} WITH PASSWORD '{database_password}' LOGIN;
                      """
                      create_role_res = execute_statement(resource_arn, database_name, secret_arn, create_role)
                      logger.info(f"Create Role Response: {create_role_res}")

                      create_schema = f"""
                      CREATE SCHEMA {schema_name};
                      """
                      create_schema_res = execute_statement(resource_arn, database_name, secret_arn, create_schema)
                      logger.info(f"Create Schema Response: {create_schema_res}")

                      grant_schema = f"""
                      GRANT ALL ON SCHEMA {schema_name} to {user_name};
                      """
                      grant_schema_res = execute_statement(resource_arn, database_name, secret_arn, grant_schema)
                      logger.info(f"Grant Schema Response: {grant_schema_res}")

                      create_table = f"""
                      CREATE TABLE {table_name} ({primary_key_field} uuid PRIMARY KEY, {vector_field} vector({dimension}), {text_field} text, {metadata_field} json)
                      """
                      create_table_res = execute_statement(resource_arn, database_name, secret_arn, create_table)
                      logger.info(f"Create Table Response: {create_table_res}")

                      grant_table = f"""
                      GRANT ALL ON TABLE {table_name} TO {user_name};
                      """
                      grant_table_res = execute_statement(resource_arn, database_name, secret_arn, grant_table)
                      logger.info(f"Grant Table Response: {grant_table_res}")

                      create_index = f"""
                      CREATE INDEX on {table_name} USING hnsw ({vector_field} vector_cosine_ops);
                      """
                      create_index_res = execute_statement(resource_arn, database_name, secret_arn, create_index)
                      logger.info(f"Create Index Response: {create_index_res}")

                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
                  if event['RequestType'] == 'Update':
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
                  if event['RequestType'] == 'Delete':
                      cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as e:
                  cfnresponse.send(event, context, cfnresponse.FAILED, {'Message': str(e)})

今回のテンプレートの肝となる箇所です。

Knowledge Base for Amazon BedrockのベクターストアとしてAuroraを利用するためには、以下の通りSQLによるベクターストアのセットアップを実施する必要があります。

こちらはCloudFormationでは対応していないため、カスタムリソースを使用してリクエストタイプがCreateとなる時のみ、セットアップ用のLambda関数を実行する事で、ベクターストアのセットアップを自動化します。

カスタムリソース及びリクエストタイプについては、以下の記事をご参照ください。

またセットアップ用のLambda関数では、Auroraに対してHTTP経由でアクセスを行える機能であるRDS Data APIを使用します。

RDS Data APIについては、以下の記事をご参照ください。

カスタムリソースとして定義したLambda関数では、以下のセットアップ処理に関するSQLをRDS Data APIを介して実行します。

• ベクターストア用の拡張機能の作成
• ユーザーの作成
• スキーマの作成
• スキーマに対する権限をユーザーに付与
• テーブルの作成
• テーブルに対する権限をユーザーに付与
• インデックスの作成

なお公式ドキュメントでは「テーブルに対する権限をユーザーに付与」処理はありませんが、ない場合Knowledge Base for Amazon Bedrock作成時にエラーが生じるため追加しています。

Knowledge Base for Amazon Bedrock

  BedrockKnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [bedrock.amazonaws.com]
            Action: ['sts:AssumeRole']
  BedrockKnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      Name: !Sub ${AWS::StackName}-knowledge-base
      KnowledgeBaseConfiguration:
          Type: VECTOR
          VectorKnowledgeBaseConfiguration:
              EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/

      RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
      StorageConfiguration:
        Type: RDS
        RdsConfiguration:
          CredentialsSecretArn: !Ref SecretForAurora
          DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
          FieldMapping:
            MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
            PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
            TextField: !FindInMap [DatabaseMap, TextField, Name]
            VectorField: !FindInMap [DatabaseMap, VectorField, Name]
          ResourceArn: !GetAtt AuroraCluster.DBClusterArn
          TableName: !FindInMap [DatabaseMap, TableName, Name]
    DependsOn: SetupAuroraData
  BedrockKnowledgeBaseDS:
    Type: AWS::Bedrock::DataSource
    Properties:
      KnowledgeBaseId: !Ref BedrockKnowledgeBase
      Name: !Sub ${AWS::StackName}-data-source
      DataSourceConfiguration:
        Type: S3
        S3Configuration:
          BucketArn: !GetAtt DataSourceBucket.Arn

最後にKnowledge Base for Amazon Bedrockを作成し、事前に作成したAurora、Secrets Manager、S3を紐づけます。

EmbeddingModelArnでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのARNを指定します。

またRdsConfigurationでは、Auroraに対する接続情報及びデータベースのスキーマ情報を指定します。

動作確認

上記のテンプレートをデプロイする事で、Knowledge Base for Amazon Bedrockを利用して正常に文書データを検索できるか確認します。

なお今回は、2024/4時点でKnowledge Base for Amazon Bedrockが利用可能なバージニアリージョンus-east-1でデプロイを行います。

CloudFormation テンプレートのデプロイ

今回は以下のパラメータで、CloudFormationスタックをデプロイします。

DatabasePasswordでは、Auroraに接続するためのパスワードを指定します。

機密情報となるため、適切なパスワードを使用してください。

EmbeddingModelIdでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのIDを指定します。

今回はamazon.titan-embed-text-v1を指定します。

なおAuroraのデプロイ時間が長く、およそ20-30分程度かかります。

S3バケットへのデータ投入

次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。

# S3 バケット名を設定

BUCKET_NAME=<バケット名>

# AWS の公式ドキュメントの PDF ファイルをダウンロード

mkdir -p pdf
cd pdf
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf -O Bedrock.pdf
cd ..

# S3 バケットに PDF ファイルをアップロード

aws s3 cp pdf s3://${BUCKET_NAME} --recursive

ファイルのアップロードが成功すると、以下の通り作成したS3バケットに4つのPDFファイルが作成されます。

Knowledge Baseデータソースの同期

続いてBedrockのコンソールに移動し、左ペインの"ナレッジベース"に切り替え、作成したナレッジベースをクリックします。

ナレッジベース画面中部あたりのデータソースに移動し、同期をクリックします。

同期中は、ステータスがSyncingになります。

データソースの同期にも、およそ20-30分程度かかります。

同期が完了すると、ステータスがReadyに戻り、最終同期時刻が表示されます。

Knowledge Baseでの検索実行

最後に検索が可能か確認するために、ナレッジベース画面のテストをクリックし、ナレッジベーステストを実施します。

例えば"回答を生成"をOFFにした状態で、「IPアドレス」と入力し送信すると、該当の用語に類似するデータソースのチャンクが表示されます。

また"回答を生成"をONにした状態で、「Lambdaに割り当てられる最大メモリを教えてください」と入力し送信すると、データソースのチャンクを元に正しく「10,240MB」という回答を返します。

これでナレッジベースが正しく動作している事を確認できました！

最後に

今回はベクターストアにAuroraを利用したKnowledge Base for Amazon Bedrockを、CloudFormation一撃で作成可能なテンプレートについて紹介しました。

今までKnowledge Baseはコンソールから作成する必要がありましたが、CloudFormationによるデプロイが可能となった事で、より簡単にKnowledge Baseを作成できるようになっています。

Knowledge Base for Amazon Bedrockをまだ触った事がない方は、ぜひ一度お試しください！

以上、つくぼし(tsukuboshi0755)でした！

参考文献

テンプレートの作成にあたって、以下のブログを参考にさせて頂きました。